和为 K 的子数组
给你一个整数数组 nums 和一个整数 k ,请你统计并返回 该数组中和为 k 的子数组的个数 。
子数组是数组中元素的连续非空序列。
示例 1:
输入:nums = [1,1,1], k = 2 输出:2 示例 2:
输入:nums = [1,2,3], k = 3 输出:2
简单做法:双指针遍历
进阶做法: 使用前缀和+哈希表的方法。这种方法能够在O(n)时间复杂度内解决问题,比暴力枚举所有子数组的O(n²)效率更高。
方法思路
- 前缀和定义:前缀和
prefix[i]表示数组前i个元素的和。 - 子数组和转化:对于区间
[i, j]的子数组,其和为prefix[j] - prefix[i-1]。 - 哈希表优化:使用哈希表记录每个前缀和出现的次数,遍历数组时,检查当前前缀和与k的差值是否存在于哈希表中,若存在则累加对应次数。
代码实现
function subarraySum(nums, k) {
const prefixCount = new Map();
prefixCount.set(0, 1); // 初始化前缀和为0的情况出现1次
let count = 0;
let currentSum = 0;
for (const num of nums) {
currentSum += num;
// 检查是否存在前缀和为currentSum - k的情况
const target = currentSum - k;
if (prefixCount.has(target)) {
count += prefixCount.get(target);
}
// 更新当前前缀和的出现次数
prefixCount.set(currentSum, (prefixCount.get(currentSum) || 0) + 1);
}
return count;
}
代码解释
- 哈希表初始化:
prefixCount存储前缀和及其出现次数,初始时前缀和0出现1次。 - 遍历数组:累加当前元素到
currentSum,计算目标值target = currentSum - k。 - 检查目标值:若哈希表中存在
target,说明存在以当前位置结尾的子数组和为k,累加其出现次数。 - 更新哈希表:将当前前缀和的出现次数更新到哈希表中。
复杂度分析
- 时间复杂度:O(n),仅需一次遍历数组。
- 空间复杂度:O(n),主要用于哈希表存储前缀和及其次数。
示例分析
以数组[1, 1, 1]和k=2为例:
- 初始状态:
currentSum = 0,prefixCount = {0: 1}。 - 遍历第一个元素1:
currentSum = 1,target = 1 - 2 = -1(不存在),更新prefixCount = {0: 1, 1: 1}。 - 遍历第二个元素1:
currentSum = 2,target = 2 - 2 = 0(存在,次数1),累加1,更新prefixCount = {0: 1, 1: 1, 2: 1}。 - 遍历第三个元素1:
currentSum = 3,target = 3 - 2 = 1(存在,次数1),累加1,更新prefixCount = {0: 1, 1: 1, 2: 1, 3: 1}。 - 结果:共有2个子数组和为2(
[1, 1]和[1, 1])。
其他解法对比
- 暴力枚举:遍历所有子数组,时间复杂度O(n²)。
- 前缀和数组:预处理前缀和数组,每次查询仍需O(n),总时间O(n²)。
- 哈希表优化:本文方法,时间O(n),空间O(n),最优解。
一、前缀和的基本概念
1. 什么是前缀和?
前缀和是一个数组 prefix,其中 prefix[i] 表示数组前 i 个元素的和。例如,对于数组 nums = [1, 2, 3],其前缀和数组为:
prefix[0] = 0 (空数组的和)
prefix[1] = 1 (nums[0])
prefix[2] = 1+2=3 (nums[0]+nums[1])
prefix[3] = 1+2+3=6 (nums[0]+nums[1]+nums[2])
2. 子数组和与前缀和的关系
对于任意子数组 nums[i...j](从索引 i 到 j,包含两端),其和可以表示为:
nums[i] + nums[i+1] + ... + nums[j] = prefix[j+1] - prefix[i]
示例:子数组 nums[1...2] = [2, 3] 的和为 2+3=5,而 prefix[3] - prefix[1] = 6 - 1 = 5,完全一致。
二、前缀和解法的核心逻辑
假设我们需要找到和为 k 的子数组,即满足:
prefix[j+1] - prefix[i] = k
变形得:
prefix[i] = prefix[j+1] - k
核心思想:
- 遍历数组时,维护当前前缀和
currentSum(即prefix[j+1])。 - 对于每个
currentSum,查询之前是否存在前缀和为currentSum - k的情况。 - 若存在,则说明存在一个子数组
nums[i...j]的和为k。
三、哈希表的作用:记录前缀和的出现次数
为什么需要记录次数?因为可能存在多个不同的 i 对应相同的 prefix[i]。例如:
- 数组
[1, 0, 1],前缀和为[0, 1, 1, 2]。 - 当
currentSum = 2,k = 1时,currentSum - k = 1,而前缀和1出现了2次(索引1和2),因此存在2个子数组和为1:[1,0]和[0,1]。
哈希表的操作步骤:
- 初始化时存入
{0: 1},因为前缀和0对应空数组,用于处理子数组从第一个元素开始的情况(如nums[0...j]的和为k时,prefix[j+1] - 0 = k)。 - 遍历数组,每次计算
currentSum,并查询currentSum - k是否在哈希表中。 - 若存在,将对应次数累加到结果中。
- 将当前
currentSum的出现次数存入哈希表(次数+1)。
四、示例解析:nums = [1, 1, 1], k = 2
步骤分解:
- 初始化:
prefixCount = {0: 1},count = 0,currentSum = 0。 - 遍历第一个元素 1:
currentSum = 0 + 1 = 1target = 1 - 2 = -1,哈希表中不存在,count不变。- 哈希表更新:
{0: 1, 1: 1}。
- 遍历第二个元素 1:
currentSum = 1 + 1 = 2target = 2 - 2 = 0,哈希表中存在0: 1,count += 1(此时count = 1)。- 哈希表更新:
{0: 1, 1: 1, 2: 1}。
- 遍历第三个元素 1:
currentSum = 2 + 1 = 3target = 3 - 2 = 1,哈希表中存在1: 1,count += 1(此时count = 2)。- 哈希表更新:
{0: 1, 1: 1, 2: 1, 3: 1}。
- 最终结果:
count = 2,对应子数组[1,1](索引0-1)和[1,1](索引1-2)。
五、为什么这种方法是 O(n) 时间复杂度?
- 每个元素仅被访问一次,计算
currentSum和查询哈希表的操作均为 O(1)。 - 哈希表的插入和查询操作平均时间复杂度为 O(1)。
- 无需像暴力解法那样嵌套遍历所有子数组(O(n²))。
六、常见疑问解答
-
为什么前缀和数组的初始值是 0?
因为子数组可能从第一个元素开始(如nums[0...j]),此时i=0,对应prefix[0] = 0。 -
如果数组中有负数,这种方法还适用吗?
适用。前缀和可能增大或减小,但哈希表仍能正确记录其出现次数。 -
如何处理重复的前缀和?
哈希表记录的是次数,例如前缀和5出现3次,说明有3个不同的i使得prefix[i] = 5,每个i都可能对应一个和为k的子数组。
七、总结:前缀和解法的三步曲
- 转化问题:将「子数组和为k」转化为「前缀和差值为k」。
- 哈希表加速:用哈希表记录前缀和及其出现次数,避免重复计算。
- 动态累加:遍历过程中实时查询和更新,确保线性时间复杂度。
通过这种方法,我们将看似需要双重循环的问题优化到了线性时间,这是算法中「空间换时间」思想的典型应用。如果还有具体步骤不理解,可以告诉我,我会进一步拆解!